package com.sunlake.spring.main.common.utils;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * Utility to parse RSA keys from PEM strings or files.
 */
public final class RsaKeyUtil {
    private RsaKeyUtil() {}

    public static PrivateKey parsePrivateKeyPem(String pem) throws NoSuchAlgorithmException, InvalidKeySpecException {
        String trimmed = pem.trim();
        byte[] pkcs8Bytes;
        if (trimmed.contains("-----BEGIN RSA PRIVATE KEY-----")) {
            String base64 = stripPemHeadersExact(trimmed, "RSA PRIVATE KEY");
            byte[] pkcs1 = Base64.getMimeDecoder().decode(base64);
            pkcs8Bytes = convertPkcs1ToPkcs8(pkcs1);
        } else {
            // default expect PKCS#8
            String base64 = stripPemHeadersExact(trimmed, "PRIVATE KEY");
            pkcs8Bytes = Base64.getMimeDecoder().decode(base64);
        }
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(pkcs8Bytes);
        return KeyFactory.getInstance("RSA").generatePrivate(spec);
    }

    public static PublicKey parsePublicKeyPem(String pem) throws NoSuchAlgorithmException, InvalidKeySpecException {
        String normalized = stripPemHeadersExact(pem, "PUBLIC KEY");
        byte[] der = Base64.getMimeDecoder().decode(normalized);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(der);
        return KeyFactory.getInstance("RSA").generatePublic(spec);
    }

    public static String readAll(Path path) throws IOException {
        return Files.readString(path, StandardCharsets.UTF_8);
    }

    /**
     * Read content from a location:
     * - classpath:/path/to/file.pem
     * - file:/abs/path.pem
     * - plain path (absolute or relative)
     */
    public static String readLocation(String location) throws IOException {
        if (location == null || location.isBlank()) {
            return null;
        }
        if (location.startsWith("classpath:")) {
            String p = location.substring("classpath:".length());
            if (!p.startsWith("/")) {
                p = "/" + p;
            }
            try (InputStream is = RsaKeyUtil.class.getResourceAsStream(p)) {
                if (is == null) {
                    throw new IOException("Classpath resource not found: " + location);
                }
                return new String(is.readAllBytes(), StandardCharsets.UTF_8);
            }
        }
        if (location.startsWith("file:")) {
            return Files.readString(Path.of(location.substring("file:".length())), StandardCharsets.UTF_8);
        }
        return Files.readString(Path.of(location), StandardCharsets.UTF_8);
    }

    private static String stripPemHeadersExact(String pem, String type) {
        String begin = "-----BEGIN " + type + "-----";
        String end = "-----END " + type + "-----";
        String body = pem.replace(begin, "").replace(end, "");
        return body.replaceAll("\\s", "").trim();
    }

    private static byte[] convertPkcs1ToPkcs8(byte[] pkcs1) {
        try {
            // PKCS#8: SEQUENCE(
            //   INTEGER 0,
            //   SEQUENCE( OID 1.2.840.113549.1.1.1, NULL ),
            //   OCTET STRING (PKCS#1 bytes)
            // )
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            ByteArrayOutputStream seqAlgId = new ByteArrayOutputStream();
            // AlgorithmIdentifier = SEQUENCE( OID rsaEncryption, NULL )
            seqAlgId.write(0x30);
            byte[] oid = new byte[]{0x06, 0x09, 0x2A, (byte)0x86, 0x48, (byte)0x86, (byte)0xF7, 0x0D, 0x01, 0x01, 0x01};
            ByteArrayOutputStream algContent = new ByteArrayOutputStream();
            algContent.write(oid);
            algContent.write(new byte[]{0x05, 0x00}); // NULL
            byte[] algBytes = algContent.toByteArray();
            seqAlgId.write(encodeLength(algBytes.length));
            seqAlgId.write(algBytes);
            byte[] algId = seqAlgId.toByteArray();

            // PrivateKeyInfo content
            ByteArrayOutputStream pkiContent = new ByteArrayOutputStream();
            // version INTEGER 0
            pkiContent.write(new byte[]{0x02, 0x01, 0x00});
            // algorithm identifier
            pkiContent.write(algId);
            // privateKey OCTET STRING
            pkiContent.write(0x04);
            pkiContent.write(encodeLength(pkcs1.length));
            pkiContent.write(pkcs1);
            byte[] pki = pkiContent.toByteArray();

            // Wrap with SEQUENCE
            out.write(0x30);
            out.write(encodeLength(pki.length));
            out.write(pki);
            return out.toByteArray();
        } catch (IOException e) {
            throw new IllegalArgumentException("Failed to convert PKCS#1 to PKCS#8: " + e.getMessage(), e);
        }
    }

    private static byte[] encodeLength(int length) {
        if (length < 0x80) {
            return new byte[]{(byte) length};
        }
        int tmp = length;
        int numBytes = 0;
        while (tmp > 0) {
            tmp >>= 8;
            numBytes++;
        }
        byte[] result = new byte[1 + numBytes];
        result[0] = (byte) (0x80 | numBytes);
        for (int i = numBytes; i > 0; i--) {
            result[i] = (byte) (length & 0xFF);
            length >>= 8;
        }
        return result;
    }
}
